from numpy import array, concatenate, linspace, zeros
import png2obj.mesh
import png2obj.image

class Heightmap():
    def __init__(self, data):
        self.data = data

    def getSize(self):
        return self.data.shape

    def toMesh(self, name):
        result = mesh.Mesh(name)
        
        nx, ny = self.getSize()
        borderedData = zeros((nx + 2, ny + 2))
        borderedData[1:-1, 1:-1] = self.data
        borderedCoordsX = linspace(-0.5*(nx+1), 0.5*(nx+1), nx+2)
        borderedCoordsY = linspace(-0.5*(ny+1), 0.5*(ny+1), ny+2)

        squareOffsetX = [0, 1, 1, 0, 0]
        squareOffsetY = [0, 0, 1, 1, 0]

        triangles = []
        
        for ix in range(nx+1):
            x = borderedCoordsX[ix]
            for iy in range(ny+1):
                y = borderedCoordsY[iy]
                square = borderedData[ix:ix+2, iy:iy+2]
                xc = x + 0.5
                yc = y + 0.5
                zc = square.mean()
                vertexCP = array([[xc, yc, zc]]).T;
                vertexCN = array([[xc, yc, -zc]]).T;
                texCoordC = array([[0.5, 0.5]]).T;
                for it in range(4):
                    x0 = x + squareOffsetX[it]
                    x1 = x + squareOffsetX[it+1]
                    y0 = y + squareOffsetY[it]
                    y1 = y + squareOffsetY[it+1]
                    z0 = square[squareOffsetX[it], squareOffsetY[it]]
                    z1 = square[squareOffsetX[it+1], squareOffsetY[it+1]]
                    vertex0p = array([[x0, y0, z0]]).T;
                    vertex1p = array([[x1, y1, z1]]).T;
                    vertex0n = array([[x0, y0, -z0]]).T;
                    vertex1n = array([[x1, y1, -z1]]).T;
                    texCoord0 = array([[squareOffsetX[it], squareOffsetY[it]]]).T;
                    texCoord1 = array([[squareOffsetX[it+1], squareOffsetY[it+1]]]).T;
                    if z0 == 0.0 and z1 == 0.0 and zc == 0.0:
                        continue
                    faceP = concatenate((vertexCP, vertex0p, vertex1p), 1)
                    faceN = concatenate((vertexCN, vertex1n, vertex0n), 1)
                    texCoordsP = concatenate((texCoordC, texCoord0, texCoord1), 1)
                    texCoordsN = concatenate((texCoordC, texCoord1, texCoord0), 1)
                    result.addFace(faceP, texCoordsP)
                    result.addFace(faceN, texCoordsN)        
        return result

if __name__ == "__main__":
    x = image.loadImage('test.png')
    hm = Heightmap(x[3, :, :] / 255.0) 
    m = hm.toMesh("test")
    s = m.toObj()
    print(s)
    outfile = open('test.obj', 'w')
    outfile.write(s)
    outfile.close()
                
                
